from datetime import datetime

import pandas as pd
import numpy as np
import xarray as xr

import ttide
from ttide.t_tide import t_tide
from scipy.optimize import curve_fit


from ppgnss import gnss_time
from ppgnss import gnss_utils
from ppgnss import gnss_geodesy

from plot_baseline import read_seal_file, read_gamit_solution
import matplotlib.pyplot as plt

# 定义周期信号模型
def periodic_signal(t, A, phi):
    # 23小时55分55秒的周期对应的角频率
    omega = 2 * np.pi / (23 * 3600 + 55 * 60 + 55)
    return A * np.sin(omega * t + phi)

gamit_solution = "../baseline/solution.dat"

zihe_excel_file_path_02 = "../data/cczihe02.xlsx"
zihe_excel_file_path_03 = "../data/cczihe03.xlsx"
excel_file_path = zihe_excel_file_path_02

seal_down_file = "../data/sea_levels/downstream_7308.txt"
seal_up_file = "../data/sea_levels/upstream_3020.txt"

lat1, lon1 = 30.22863399, 120.73933035
lat2, lon2 = 30.22401975, 120.74423075

x1, y1, z1 = gnss_geodesy.blh2xyz(lat1, lon1, 0)
x2, y2, z2 = gnss_geodesy.blh2xyz(lat2, lon2, 0)

north, earth, up = gnss_geodesy.dxyz2neu([x2-x1, y2-y1, z2-z1], [x1, y1, z1])
alpha = np.arctan(north/earth)

df = pd.read_excel(excel_file_path)
xr_gg = read_gamit_solution(gamit_solution, 2022, 118)
xr_seal_down = read_seal_file(seal_down_file)

xr_seal_up = read_seal_file(seal_up_file)

# ['采集时间', '设备序列号', '东向坐标 (m)', '北向坐标 (m)', '天向坐标 (m)', 
# '东向坐标日变形速率 (mm/d)',  '北向坐标日变形速率 (mm/d)', '天向坐标日变形速率 (mm/d)',
# '平面日变形速率 (mm/d)', '三维日变形速率 (mm/d)', '差分期龄 (s)']
count = len(df["采集时间"])
time_from = np.datetime64(df["采集时间"][count-1])
time_to = np.datetime64(df["采集时间"][0])
dtime = np.datetime64(df["采集时间"][count-2])-np.datetime64(df["采集时间"][count-1])
nepoch = (time_to - time_from)//dtime+1

times = np.array([time_from + dtime*_ - 7.5*np.timedelta64(1, "h") for _ in range(nepoch)])

dus = df["天向坐标 (m)"]-df["天向坐标 (m)"][0]
dns = df["北向坐标 (m)"]-df["北向坐标 (m)"][0]
des = df["东向坐标 (m)"]-df["东向坐标 (m)"][0]
err_max = 0.04 # 0.04
idx = np.logical_or(np.logical_or(np.abs(dus) > err_max, np.abs(dns) > err_max), np.abs(des)> err_max)
dus[idx] = np.nan
dns[idx] = np.nan
des[idx] = np.nan

shape = (len(times), 5)
data = np.full(shape, np.nan)

xr_data = xr.DataArray(data, dims=["time", "data"], coords=[times, ['n', 'e', 'u', "v", "a"]])
valid_times = np.array([np.datetime64(_) - 7.5*np.timedelta64(1, "h") for _ in df["采集时间"]]) 

v = des*np.cos(alpha) + dns*np.sin(alpha)
a = -des*np.sin(alpha) + dns*np.cos(alpha)
# import sys
# sys.exit(0)

xr_data.loc[valid_times, "u"] = dus
xr_data.loc[valid_times, "n"] = dns
xr_data.loc[valid_times, "e"] = des
xr_data.loc[valid_times, "v"] = v
xr_data.loc[valid_times, "a"] = a


# xr_time1, xr_time2 = xr.align(xr_data.coords["time"], xr_seal_down.coords["time"])
# print(xr_time1)
time_resample = [np.datetime64(str(times[0])[:10]+" 00:00:00")+np.timedelta64(_, "h") for _ in range(500*24)]

# xr_data, xr_seal_down = xr.align(xr_data, xr_seal_down, join="left")
# xr_data, xr_seal_up = xr.align(xr_data, xr_seal_up, join="left")
xr_seal_delta = xr_seal_up - xr_seal_down

time_spans = [# ("2022-05-20 10:00:00", "2022-05-28 16:00:00"),  # 8 days
              # ("2022-05-30 08:00:00", "2022-06-13 00:00:00"),  # 13 days
              # ("2022-05-20 10:00:00", "2022-06-13 00:00:00"), # 23 days
              # ("2022-05-20 10:00:00", "2022-09-13 00:00:00"), # 23 days
              # ("2022-05-20 10:00:00", "2023-03-09 00:00:00"), # 23 days
              # ("2022-10-21", "2022-10-29 08:00:00"), # 8 days
              # ("2022-10-30 08:00:00", "2022-11-13"), # 13 days
              # ("2022-10-21", "2022-11-13"), # 13 days
              # ("2023-01-03 00:00:00", "2023-01-14 00:00:00"), # 11 days
              ("2023-01-29 00:00:00", "2023-03-09 00:00:00"), # 50 days
              # ("2022-10-22 00:00:00", "2022-11-18 00:00:00"), # 50 days
              ]
freqs = [0.0015122, 0.0028219,0.0357064,0.0372185,0.0387307,0.0402686,0.0417807,0.0463430,0.0776895,0.0789992,0.0805114,0.0820236,0.0833333,0.1192421,0.1222921,0.1251141,0.1610228,0.1638447,0.1666667,0.2415342,0.2443561,0.2471781]
names = ["MM","MSF","2Q1","Q1","O1","NO1","K1","UPS1","MU2","N2","M2","L2","S2","MO3","MK3","SK3","M4","MS4","S4","M6","2MS6","2SM6"]
freqs=[0.0015122, 0.0028219, 0.0357064, 0.0372185, 0.0387307, 0.0402686, 0.0417807, 0.0776895, 0.0789992, 0.0805114, 0.0820236, 0.0833333, 0.1192421, 0.1222921, 0.1251141, 0.1610228, 0.1638447, 0.1666667, 0.2084474, 0.2415342, 0.2443561]
names = ["MM", "MSF", "2Q1", "Q1", "O1", "NO1", "K1", "MU2", "N2", "M2", "L2", "S2", "MO3", "MK3", "SK3", "M4", "MS4", "S4", "2SK5", "M6", "2MS6"]
freq=[0.0015122,0.0028219,0.0357064,0.0372185,0.0387307,0.0402686,0.0417807,0.0776895,0.0789992,0.0805114,0.0820236,0.0833333,0.1192421,0.1222921,0.1251141,0.1610228,0.1638447,0.1666667,0.2028035,0.2415342,0.2443561,0.3220456]
names=["MM","MSF","2Q1","Q1","O1","NO1","K1","MU2","N2","M2","L2","S2","MO3","MK3","SK3","M4","MS4","S4","2MK5","M6","2MS6","M8"]
y_lim_min, y_lim_max = -0.02, 0.02
for time_from, time_to in time_spans:
    fig, axes = plt.subplots(nrows=4, ncols=2, sharey=False, sharex="col", figsize=(8, 8))
    xr_neu = xr_data.loc[time_from:time_to]
    xr_seal_down = xr_seal_down.loc[time_from:time_to]
    xr_seal_up = xr_seal_up.loc[time_from:time_to]
    xr_delta = xr_seal_up - xr_seal_down
    dirs = ["u", "e", "n"]
    dirs = ["u", "v", "a"]
    for icom, com in enumerate(dirs):
        obs1 = xr_neu.loc[:, com]
        stime = datetime.strptime(str(obs1.coords["time"].values[0])[:19], '%Y-%m-%dT%H:%M:%S')
        print(stime)
        obs1 = obs1 - obs1.mean()
        # obs1 = obs1*10
        # obs1 = obs1.interpolate_na(dim="time")
        out1 = t_tide(obs1.values, dt=5/60., stime=stime, lat=np.array(30.23), synth=1)

        axes[icom][0].plot(obs1.coords["time"], obs1, linestyle='-', alpha=0.7, label=com.upper())
        axes[icom][0].set_ylim((y_lim_min, y_lim_max))
        axes[icom][0].plot(obs1.coords["time"], out1["xout"].squeeze(), linestyle='--', linewidth=1, label=u'Prediction')
        # axes[icom][0].plot(obs1.coords["time"], obs1-out1["xout"].squeeze(), label=u'Error')
        # axes[icom][0].fill_between(obs1.coords["time"], obs1-out1["xout"].squeeze(), where=(obs1-out1["xout"].squeeze()< 0), color="gray", alpha=0.5)
        # axes[icom][0].fill_between(obs1.coords["time"], obs1-out1["xout"].squeeze(), where=(obs1-out1["xout"].squeeze()> 0), color="gray", alpha=0.5)


        axes[icom][0].legend(numpoints=1, loc='lower right')
        axes[icom][0].grid(axis='y', linestyle='--')
        axes[icom][0].set_yticks(np.arange(-0.02, 0.03, 0.005))
        fs = 60/5  # 采样频率，1小时采样一次
        signal = obs1.values
        signal = np.nan_to_num(signal, copy=True, nan=0)
        fft_result = np.fft.fft(signal)
        fft_freqs = np.fft.fftfreq(len(fft_result), 1/fs)  # 计算频率轴
    # 提取正频率分量
        positive_freq_mask = fft_freqs > 0
        fft_freqs_positive = fft_freqs[positive_freq_mask]
        fft_result_positive = fft_result[positive_freq_mask] 
        axes[icom][1].plot(fft_freqs_positive, np.abs(fft_result_positive))
        axes[icom][1].set_ylim((0, np.nanmax(np.abs(fft_result_positive))))

        for freq, name in zip(freqs, names):
            axes[icom][1].axvline(x=freq, color='red', linestyle='-', linewidth=0.5, label=name)
            axes[icom][1].text(freq, 0.2, name, fontsize=10)


    # obs = xr_delta - xr_delta.mean()
    obs = xr_seal_down - xr_seal_down.mean()
    stime = datetime.strptime(str(obs.coords["time"].values[0])[:19], '%Y-%m-%dT%H:%M:%S')
    out = t_tide(obs, dt=5/60, stime=stime, lat=np.array(30.23), synth=1)
    axes[3][0].plot(obs.coords["time"], obs, label=u'Sea Level')
    # obs2 = xr_seal_up - xr_seal_up.mean()
    axes[3][0].plot(obs.coords["time"], out["xout"].squeeze(), label=u'Prediction')
    axes[3][0].set_ylim((-5, 5))

    # axes[3].plot(xr_seal_up.coords["time"], xr_seal_up, label=u'Up Sea Level', c="r")
    # axes[3].plot(xr_seal_down.coords["time"], xr_seal_down, label=u'Down Sea Level', c="k")
    axes[3][0].legend(numpoints=1, loc='lower right')

    fs = 60/5  # 采样频率，1小时采样一次
    signal = obs.values
    signal = np.nan_to_num(signal, nan=0)
    fft_result = np.fft.fft(signal)
    fft_freqs = np.fft.fftfreq(len(fft_result), 1/fs)  # 计算频率轴
    # 提取正频率分量
    positive_freq_mask = fft_freqs > 0
    fft_freqs_positive = fft_freqs[positive_freq_mask]
    fft_result_positive = fft_result[positive_freq_mask]    
    axes[3][1].plot(fft_freqs_positive, np.abs(fft_result_positive))
    for freq, name in zip(freqs, names):
        axes[3][1].axvline(x=freq, color='red', linestyle='-', linewidth=0.5, label=name)
        axes[3][1].text(freq, 100, name, fontsize=10)
    axes[3][1].set_title('Frequency Spectrum')
    axes[3][1].set_xlabel('Frequency (cycles per hour)')

    plt.show() 



